#include <opencv2/core/mat.hpp>
#include <onnxruntime_cxx_api.h>

using cv::Mat;
using std::string;
using std::vector;

struct NetConfig
{
	float confidence_threshold;
	float nms_threshold;
	string model_path;
	string classes_path;
};

typedef struct BoxInfo
{
	float x1;
	float y1;
	float x2;
	float y2;
	float score;
	int label;
} BoxInfo;

class YOLOV6
{
public:
	YOLOV6(NetConfig config);
	void detect(Mat &frame);

private:
	int inpWidth;
	int inpHeight;
	int nout;
	int num_proposal;
	vector<string> class_names;
	int num_class;

	float confidence_threshold;
	float nms_threshold;
	const bool keep_ratio = true;
	vector<float> input_image_;
	void normalize_(Mat img);
	void nms(vector<BoxInfo> &input_boxes);
	Mat resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left);

	Ort::Env env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "YOLOV6");
	Ort::Session *ort_session = nullptr;
	Ort::SessionOptions sessionOptions = Ort::SessionOptions();
	vector<string> input_names;
	vector<string> output_names;
	vector<vector<int64_t>> input_node_dims;	// >=1 outputs
	vector<vector<int64_t>> output_node_dims; // >=1 outputs
};
